import logging
import random
import time
import weakref

import gymnasium as gym
import numpy as np
import pyproj

from tcp_comm import *

logger = logging.getLogger(__name__)

wgs84 = pyproj.CRS("EPSG:4326")
geocentric = pyproj.CRS("EPSG:32651")

# Coordinate transform
transformer_wgs84 = pyproj.Transformer.from_crs(geocentric, wgs84, always_xy=True)
transformer_geocentric = pyproj.Transformer.from_crs(wgs84, geocentric, always_xy=True)
offsets = (644911.9272301428, -1769688.8120349043, 0.0)


def to_lonlat(x, y, z, offsets=offsets):
    x -= offsets[0]
    y -= offsets[1]
    z -= offsets[2]
    lon, lat, alt = transformer_wgs84.transform(x, y, z)
    return lon, lat, alt


def to_xyz(lon, lat, alt, offsets=offsets):
    x, y, z = transformer_geocentric.transform(lon, lat, alt)
    x += offsets[0]
    y += offsets[1]
    z += offsets[2]
    return x, y, z


def to_grid(x, y, z):
    pass


ACTION_MAP = {
    0: np.array([10, 0, 0]),  # Move right
    1: np.array([-10, 0, 0]),  # Move left
    2: np.array([0, 10, 0]),  # Move forward
    3: np.array([0, -10, 0]),  # Move backward
    4: np.array([0, 0, 10]),  # Move up
    5: np.array([0, 0, -10]),  # Move down
}


class PathPlanEnv(gym.Env):
    def __init__(self, config: dict = None):
        self.config = config or {}

        self.grid_map = self.config.get("grid_map", None)
        self.area_lon_min = self.config.get("area_lon_min", 120)
        self.area_lon_max = self.config.get("area_lon_max", 125)
        self.area_lat_min = self.config.get("area_lat_min", 25)
        self.area_lat_max = self.config.get("area_lat_max", 26)
        self.area_alt_min = self.config.get("area_alt_min", 0)
        self.area_alt_max = self.config.get("area_alt_max", 50)
        self.min_goal_distance = self.config.get("min_goal_distance", 100)
        self.max_steps = self.config.get("max_steps", 250)

        self.collided = False
        self.client = CommClient(self.config)

        weak_self = weakref.ref(self)
        self.client.register_handler(
            "collision",
            CollisionInfoMsg,
            lambda e: self.collision_callback(weak_self, e),
        )

        self.action_space = gym.spaces.Discrete(4)

    def reset(self, *, seed=None, options=None):
        super().reset(seed=seed, options=options)

        self._last_frame = None
        self.steps = 0
        self.collided = False

        if not self.client.connected():
            self.client.connect()

        while True:
            self.start_point = (
                random.uniform(self.area_lon_min, self.area_lon_max),
                random.uniform(self.area_lat_min, self.area_lat_max),
                random.uniform(self.area_alt_min, self.area_alt_max),
            )

            self.goal_point = (
                random.uniform(self.area_lon_min, self.area_lon_max),
                random.uniform(self.area_lat_min, self.area_lat_max),
                random.uniform(self.area_alt_min, self.area_alt_max),
            )

            start = np.array(to_xyz(*self.start_point))
            goal = np.array(to_xyz(*self.goal_point))
            dist = np.linalg.norm(start - goal)
            if dist > self.min_goal_distance:
                self._last_dist = dist
                break

        # self.start_point = (121.56850, 25.03757, 17.5595397949219)
        # self.goal_point = (121.57070, 25.03903, 30.5547332763672)

        uav_init_msg = ObjectInitActionMsg(
            timestamps=time.time(),
            objectName="uav",
            objectType=ObjectType.ObjectType_UAV_Red_DR400,
            objectCamp=0,
            objectPos=Position3D(*self.start_point),
            objectOrientation=RotationAngles(),
        )

        self.client.start_episode([uav_init_msg])
        freq_msg = SetMsgFrequencyMsg(time.time(), frequency=500)
        sim_speed_msg = SimulationSpeedInfoMsg(time.time(), speed=4)
        self.client.send_message(self.client._world_client, freq_msg)
        self.client.send_message(self.client._world_client, sim_speed_msg)

        obs = self.get_obs()
        return obs, {}

    def step(self, action):
        self.steps += 1
        movement = ACTION_MAP[action]
        new_pos = self.uav_xyz + movement
        new_lonlat = to_lonlat(*new_pos)

        self.client.send_control(
            ManeuveringActionMsg(
                timestamps=time.time(),
                objectName="uav",
                path=json.dumps(
                    [
                        ManeuveringInfo(
                            Position3D(*new_lonlat),
                            self.uav_vel,
                            self.uav_rot.yaw,
                        )
                    ],
                    option=json.OPT_SERIALIZE_NUMPY,
                ).decode("utf-8"),
            )
        )

        obs = self.get_obs()
        reward = self.get_reward()
        te = self._last_frame.get_object("uav").lifeValue < 100.0 or self.collided
        tr = self.steps >= self.max_steps
        info = {}
        return obs, reward, te, tr, info

    def get_obs(self):
        while True:
            frame = self.client.get_frame()
            if frame is None:
                continue
            if self._last_frame is None:
                self._last_frame = frame
                break
            if frame.timestamps != self._last_frame.timestamps:
                self._last_frame = frame
                break

        uav_state = frame.get_object("uav")
        uav_pos = uav_state.objectPosition
        uav_rot = uav_state.objectOrientation
        uav_vel = uav_state.objectVelocity
        uav_acc = uav_state.acceleration

        obs = np.concatenate(
            [
                self.goal_xyz,
                np.array(to_xyz(uav_pos.lon, uav_pos.lat, uav_pos.altitude)),
                np.array([uav_rot.yaw, uav_rot.pitch, uav_rot.roll]),
                np.array([uav_vel.x, uav_vel.y, uav_vel.z]),
                np.array([uav_acc.x, uav_acc.y, uav_acc.z]),
            ]
        )

        return obs

    def get_reward(self):
        reward = 0.0

        # distance reward
        dist = np.linalg.norm(self.uav_xyz - self.goal_xyz)
        reward += self._last_dist - dist
        self._last_dist = dist

        # step cost
        reward -= 1.0

        return reward

    def close(self):
        self.client.disconnect()

    @staticmethod
    def collision_callback(weak_self, event: CollisionInfoMsg):
        self: "PathPlanEnv" = weak_self()
        if not self:
            return
        if event.objectName == "uav":
            self.collided = True

    @property
    def uav_xyz(self):
        if self._last_frame is None:
            return None

        uav_state = self._last_frame.get_object("uav")
        uav_pos = uav_state.objectPosition
        return np.array(to_xyz(uav_pos.lon, uav_pos.lat, uav_pos.altitude))

    @property
    def uav_rot(self):
        if self._last_frame is None:
            return None

        uav_state = self._last_frame.get_object("uav")
        return uav_state.objectOrientation

    @property
    def uav_vel(self):
        if self._last_frame is None:
            return None

        uav_state = self._last_frame.get_object("uav")
        return uav_state.objectVelocity

    @property
    def goal_xyz(self):
        return np.array(to_xyz(*self.goal_point))


if __name__ == "__main__":
    import sys

    logging.basicConfig(
        level=logging.DEBUG,
        format='%(asctime)s %(levelname)s %(module)s %(funcName)s: "%(message)s"',
        datefmt="%d-%m-%y %H:%M:%S",
        handlers=[
            logging.StreamHandler(sys.stdout),
            logging.FileHandler("env.log", "w", encoding="utf-8"),
        ],
    )

    env = PathPlanEnv()

    try:
        for i in range(10):
            logger.info(f"Starting episode: {i + 1}")
            obs, info = env.reset()
            done = False

            while not done:
                action = env.action_space.sample()
                obs, reward, te, tr, info = env.step(action)
                done = te or tr
    finally:
        env.close()
